# -*- coding: UTF-8 -*-
"""
@Project ：video_demo
@File    ：complete_demo.py
@IDE     ：PyCharm
@Author  ：测试工程师
@Date    ：2025/3/31 15:30
@Describe：完整视频分析演示版
"""
import time
import threading
from collections import deque
from datetime import datetime
import cv2
import numpy as np
import torch
from pathlib import Path
from ultralytics import YOLO
import logging

# 日志配置
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
    ]
)
logger = logging.getLogger(__name__)

# 模型配置
MODEL_PATH = "/workspace/models/lyyj_dac11m_safe_v4/weights/best.pt"  # 测试用默认模型，实际替换为您的模型路径
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = YOLO(MODEL_PATH).to(device)
logger.info(f"模型加载完成 | 设备: {device}")

# 分类配置（示例配置，按实际修改）
CLASS_CONFIG = {
    'names': {0: '持械', 1: '火患', 2: '打斗', 3: '人员', 4: '动物'},
    'colors': {
        0: (255, 0, 0),  # 蓝色-持械
        1: (0, 0, 255),  # 红色-火患
        2: (0, 255, 255),  # 黄色-打斗
        3: (0, 255, 0),  # 绿色-人员
        4: (255, 255, 0)  # 青色-动物
    }
}

# 系统参数
CONFIG = {
    'fire_validation': {
        'enabled': True,                # 是否启用火患验证
        'gradient_threshold': 0.15,     # 亮度梯度阈值
        'flicker_frames': 5,            # 频域分析帧数
        'flicker_min_hz': 1,            # 最小闪烁频率(Hz)
        'flicker_max_hz': 10,           # 最大闪烁频率(Hz)
        'flicker_energy_ratio': 0.25,  # 频域能量占比阈值
        'high_conf_thresh': 0.7,       # 高置信度直接通过阈值
        'low_conf_thresh': 0.25,        # 低置信度直接过滤阈值
        'min_roi_size': 20             # 最小有效ROI尺寸
    },
    'frame_size': (640, 640),
    'base_conf': 0.30,  # 基础置信度（覆盖所有类别）
    'class_conf': {  # 分类置信度阈值
        0: 0.40,  # 持械
        1: 0.40,  # 火患
        2: 0.60,  # 打斗
        3: 0.50,  # 人员
        4: 0.50  # 动物
    },
    'alarm_interval': 30,  # 报警间隔(秒)
    'output_dir': Path("alarm_frames_test3"),
    'min_detections': 10,  # 最小连续检测次数
    'detection_window': 10,  # 检测时间窗口(秒)
    'frame_skip': 2,  # 帧跳过间隔
    'reconnect_interval': 3  # 重连间隔(秒)
}

# 测试视频配置
VIDEO_PATHS = {
    #"douo1.mp4": "/workspace/test/lyyj_dac11m_safe_v3/mvs/20250307/douo1.mp4",
    "chixie1.mp4": "/workspace/test/lyyj_dac11m_safe_v3/mvs/20250307/chixie1.mp4",
    "chixie2.mp4": "/workspace/test/lyyj_dac11m_safe_v3/mvs/20250307/chixie2.mp4",
    "chixie3.mp4": "/workspace/test/lyyj_dac11m_safe_v3/mvs/20250307/chixie3.mp4",
    "douo1.mp4": "/workspace/test/lyyj_dac11m_safe_v3/mvs/20250307/douo1.mp4",
    "douo2.mp4": "/workspace/test/lyyj_dac11m_safe_v3/mvs/20250307/douo2.mp4",
    "huoqing1.mp4": "/workspace/test/lyyj_dac11m_safe_v3/mvs/20250307/huoqing1.mp4",
    "huoqing2.mp4": "/workspace/test/lyyj_dac11m_safe_v3/mvs/20250307/huoqing2.mp4",
    "juji1.mp4": "/workspace/test/lyyj_dac11m_safe_v3/mvs/20250307/juji1.mp4",
    "juji2.mp4": "/workspace/test/lyyj_dac11m_safe_v3/mvs/20250307/juji2.mp4",
}


class VideoProcessor:
    def __init__(self, video_path, area_id):
        self.area_id = area_id
        self.video_path = video_path
        self.cap = None
        self.active = True
        self.last_alarms = {}
        self.frame_count = 0

        self.last_frames = deque(maxlen=CONFIG['fire_validation']['flicker_frames'])
        self.fire_validator = FireValidator()

        # 初始化操作
        self.connect()
        self._warmup_model()
        CONFIG['output_dir'].mkdir(exist_ok=True)

        # 启动处理线程
        threading.Thread(target=self.process, daemon=True).start()

    def _warmup_model(self):
        """模型预热"""
        warmup_frame = np.zeros((*CONFIG['frame_size'], 3), dtype=np.uint8)
        try:
            _ = model(warmup_frame, verbose=False)
            logger.info(f"[{self.area_id}] 模型预热成功 | 输入尺寸: {CONFIG['frame_size']}")
        except Exception as e:
            logger.error(f"[{self.area_id}] 模型预热失败: {str(e)}")

    def connect(self):
        """连接视频源"""
        try:
            self.cap = cv2.VideoCapture(self.video_path)
            if not self.cap.isOpened():
                raise IOError(f"无法打开视频文件: {self.video_path}")

            # 设置视频参数
            self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, CONFIG['frame_size'][0])
            self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, CONFIG['frame_size'][1])
            logger.info(f"[{self.area_id}] 视频源初始化成功 | 路径: {self.video_path}")
            return True
        except Exception as e:
            logger.error(f"[{self.area_id}] 视频源连接失败: {str(e)}")
            return False

    def process(self):
        """主处理流程"""
        while self.active and self.cap.isOpened():
            try:
                ret, frame = self.cap.read()
                if not ret:
                    logger.warning(f"[{self.area_id}] 视频读取结束或错误")
                    break

                # 维护帧缓存
                resized_frame = cv2.resize(frame, CONFIG['frame_size'])


                #self.last_frames.append(resized_frame.copy())


                # 帧跳过控制
                self.frame_count += 1
                if self.frame_count % CONFIG['frame_skip'] != 0:
                    continue

                # 执行检测
                results = self.detect(resized_frame)
                if results:
                    processed = self.post_process(results, resized_frame)
                    if processed:
                        self.check_alarms(processed, resized_frame.copy())

            except Exception as e:
                logger.error(f"[{self.area_id}] 处理异常: {str(e)}", exc_info=True)
                time.sleep(1)

        self.release()

    def detect(self, frame):
        """目标检测核心方法"""
        try:
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            results = model(
                rgb_frame,
                imgsz=CONFIG['frame_size'],
                conf=CONFIG['base_conf'],
                device=device,
                iou=1.0,  # 禁用内置NMS
                verbose=False
            )
            return [{
                'xyxy': box.xyxy[0].cpu().numpy(),
                'conf': box.conf.item(),
                'cls_id': int(box.cls.item()),
                'time': time.time()
            } for box in results[0].boxes]
        except Exception as e:
            logger.error(f"[{self.area_id}] 检测失败: {str(e)}")
            return []

    def post_process(self, detections, frame):
        """检测结果后处理"""
        valid_boxes = []

        # light = self._get_light_condition(self.last_frame)

        # 连续三个检测都误报直接结束循环
        validate_weapon_count: int = 0
        validate_fire_count: int = 0
        validate_fight_count: int = 0
        for box in detections:
            # 保存原始检测信息
            raw_cls = box['cls_id']
            raw_conf = box['conf']

            cls_threshold = CONFIG['class_conf'].get(raw_cls, CONFIG['base_conf'])

            if raw_conf < cls_threshold:
                continue

            # 持械类特殊处理
            if raw_cls == 0:

                if self.validate_weapon(box):
                    logger.debug(f"持械误报过滤")

                    validate_weapon_count += 1
                    if validate_weapon_count >= 3:
                        return valid_boxes

                    continue

            elif raw_cls == 1:  # 火患类误报处理
                if not self.validate_fire(box, frame):

                    validate_fire_count += 1
                    if validate_fire_count >= 3:
                        return valid_boxes

                    continue

            elif raw_cls == 2:  # 打斗类别特殊处理
                if not self.validate_fight(box):
                    box['cls_id'] = 3  # 转为人员类别
                    raw_cls = box['cls_id']
                    # box['conf'] = raw_conf * 0.8  # 降低置信度

            # 应用分类置信度过滤
            if box['conf'] >= cls_threshold:
                valid_boxes.append(box)

        return valid_boxes

    def validate_fire(self, box, frame):
        """火患验证整合方法"""
        if not CONFIG['fire_validation']['enabled']:
            return True

        conf = box['conf']
        cfg = CONFIG['fire_validation']

        # 置信度分级处理
        if conf > cfg['high_conf_thresh']:
            return True
        if conf < cfg['low_conf_thresh']:
            return False

        # 提取当前ROI
        x1, y1, x2, y2 = map(int, box['xyxy'])
        current_roi = frame[y1:y2, x1:x2]

        # 检查ROI有效性
        if current_roi.size == 0 or \
                (x2 - x1) < cfg['min_roi_size'] or \
                (y2 - y1) < cfg['min_roi_size']:
            return False

        # 获取历史ROI序列
        roi_sequence = []
        for f in list(self.last_frames)[-cfg['flicker_frames']:]:
            roi_sequence.append(f[y1:y2, x1:x2])

        # 执行验证
        is_flickering = self.fire_validator.check_flicker(roi_sequence)
        has_gradient = self.fire_validator.check_gradient(current_roi)

        logger.debug(f"火患验证 | 闪烁:{is_flickering} 梯度:{has_gradient} 置信度:{conf:.2f}")
        # 必须满足：无车灯特征 并且 有火焰特征
        return (not is_flickering) and has_gradient

    def validate_weapon(self, box):
        x1, y1, x2, y2 = map(int, box['xyxy'])
        w, h = x2 - x1, y2 - y1

        # 条件1：排除固定坐标误报
        if (1400 <= x1 <= 1500) and (2100 <= x2 <= 2200) and (100 <= y1 <= 130) and (1270 <= y2 <= 1280):
            return True

        # （夜间提升至0.5）
        min_conf = 0.5 if self.is_night_time() else CONFIG['class_conf'].get(box['cls_id'], CONFIG['base_conf'])
        if box['conf'] < min_conf:
            logger.debug(
                f"持械低置信度过滤({'夜间' if self.is_night_time() else '日间'}): {box['conf']:.2f} < {min_conf}")
            return True

        # 条件2：通用尺寸限制
        return ((550 <= w <= 800) and (1000 <= h <= 1400))  # or (self.is_near_border(box['xyxy']))  # 边界检查

    def is_night_time(self):
        """判断当前是否处于夜间时段（22:00-07:00）"""
        current_hour = datetime.now().hour
        return 18 <= current_hour or current_hour < 7

    def validate_fight(self, box):
        """验证打斗检测有效性"""
        x1, y1, x2, y2 = box['xyxy']
        w, h = x2 - x1, y2 - y1
        img_w, img_h = CONFIG['frame_size']

        # 基础条件检查
        conditions = [
            # box['conf'] > CONFIG['fighting_conf'],  # 置信度要求
            h > 0,  # 最小高度
            0.5 < w / h < 1.5,  # 宽高比范围
            not self.is_near_border(box['xyxy'])  # 边界检查
        ]

        return all(conditions)

    def is_near_border(self, xyxy, thickness=50):
        """检查是否靠近边界"""
        x1, y1, x2, y2 = xyxy
        img_w, img_h = CONFIG['frame_size']
        return (x1 < thickness or (img_w - x2) < thickness or
                y1 < thickness or (img_h - y2) < thickness)

    def check_alarms(self, detections, frame):
        """检查报警条件
        参数:
            detections: 检测到的目标列表
            frame: 当前视频帧
        """
        current_time = time.time()
        alarms_to_send = {}

        logger.info(f"开始检查报警条件 | 区域: {self.area_id} | 检测到 {len(detections)} 个目标")

        # 先统计当前帧中出现的类别（去重）
        detected_classes = set()
        for box in detections:
            cls_id = box['cls_id']
            cls_name = CLASS_CONFIG['names'][cls_id]
            conf = box['conf']
            detected_classes.add(cls_name)

            logger.debug(
                f"检测目标 | 类别: {cls_name}({cls_id}) | "
                f"置信度: {conf:.2f} | 坐标: {box['xyxy']}"
            )

        # 对当前帧中出现的每个类别进行计数（每帧每类别只计一次）
        for cls_name in detected_classes:
            # 初始化或获取检测记录
            if cls_name not in self.last_alarms:
                self.last_alarms[cls_name] = {'times': [], 'last_sent': 0}
                logger.debug(f"初始化报警记录 | 类别: {cls_name}")

            records = self.last_alarms[cls_name]
            records['times'].append(current_time)

            # 清理过期记录
            valid_times = [t for t in records['times']
                           if current_time - t <= CONFIG['detection_window']]
            expired_count = len(records['times']) - len(valid_times)
            records['times'] = valid_times

            if expired_count > 0:
                logger.debug(
                    f"清理过期记录 | 类别: {cls_name} | "
                    f"清理数量: {expired_count} | 剩余有效: {len(valid_times)}"
                )

            # 检查报警条件
            detections_count = len(records['times'])
            time_since_last = current_time - records['last_sent']

            logger.debug(
                f"检查报警条件 | 类别: {cls_name} | "
                f"当前检测次数: {detections_count}/{CONFIG['min_detections']} | "
                f"距上次报警: {time_since_last:.1f}s/{CONFIG['alarm_interval']}s"
            )

            if (detections_count >= CONFIG['min_detections'] and
                    time_since_last >= CONFIG['alarm_interval']):

                # 获取该类别的最高置信度
                max_conf = max(b['conf'] for b in detections if CLASS_CONFIG['names'][b['cls_id']] == cls_name)

                logger.info(
                    f"满足报警条件 | 类别: {cls_name} | "
                    f"连续检测: {detections_count}次 | "
                    f"最高置信度: {max_conf:.2f}"
                )

                # 保存报警帧
                frame_path = self.save_alarm_frame(frame, detections)
                if not frame_path:
                    logger.error(f"报警帧保存失败 | 类别: {cls_name}")
                    continue

                # 记录要发送的报警
                alarms_to_send[cls_name] = {
                    'max_conf': max_conf,
                    'timestamp': current_time,
                    'frame_path': str(frame_path)
                }

                # 更新发送时间
                records['last_sent'] = current_time
                records['times'].clear()

                logger.info(
                    f"准备发送报警 | 类别: {cls_name} | "
                    f"最后报警时间更新: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(current_time))}"
                )

        # 输出报警信息
        if alarms_to_send:
            self.print_alarms(alarms_to_send)

    def print_alarms(self, alarms):
        """打印报警信息（替代原异步发送）"""
        for cls_name, info in alarms.items():
            alarm_msg = (
                f"\n{'!' * 20} 报警触发 {'!' * 20}\n"
                f"区域: {self.area_id}\n"
                f"类别: {cls_name}\n"
                f"置信度: {info['max_conf']:.2f}\n"
                f"时间: {datetime.fromtimestamp(info['timestamp']).strftime('%Y-%m-%d %H:%M:%S')}\n"
                f"截图: {info['frame_path']}\n"
                f"{'!' * 50}"
            )
            logger.info(alarm_msg)

    def save_alarm_frame(self, frame, detections):
        """保存报警截图"""
        try:
            # 绘制检测框
            for box in detections:
                x1, y1, x2, y2 = map(int, box['xyxy'])
                color = CLASS_CONFIG['colors'][box['cls_id']]
                cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)

                # 添加标签
                label = f"{CLASS_CONFIG['names'][box['cls_id']]} {box['conf']:.2f}"
                cv2.putText(
                    frame, label, (x1, y1 - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2
                )

            # 保存文件
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
            save_path = CONFIG['output_dir'] / f"{self.area_id}_{timestamp}.jpg"
            cv2.imwrite(str(save_path), frame)
            return save_path
        except Exception as e:
            logger.error(f"[{self.area_id}] 保存报警帧失败: {str(e)}")
            return None

    def release(self):
        """释放资源"""
        self.active = False
        if self.cap and self.cap.isOpened():
            self.cap.release()
        logger.info(f"[{self.area_id}] 资源已释放")


class FireValidator:
    def __init__(self):
        cv2.ocl.setUseOpenCL(True)

    def check_gradient(self, roi):
        """亮度梯度验证"""
        if roi.size == 0:
            return False

        try:
            gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY).astype(float) / 255
            grad_x = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
            grad_y = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
            gradient = np.sqrt(grad_x ** 2 + grad_y ** 2)
            return np.mean(gradient) > CONFIG['fire_validation']['gradient_threshold']
        except Exception as e:
            logger.error(f"梯度计算失败: {str(e)}")
            return False

    def check_flicker(self, roi_sequence):
        """频域闪烁分析"""
        if len(roi_sequence) < CONFIG['fire_validation']['flicker_frames']:
            return False

        brightness = []
        for roi in roi_sequence:
            if roi.size == 0:
                return False
            try:
                h, w = roi.shape[:2]
                center_roi = roi[max(0, h // 2 - 5):min(h, h // 2 + 5), max(0, w // 2 - 5):min(w, w // 2 + 5)]
                brightness.append(np.mean(cv2.cvtColor(center_roi, cv2.COLOR_BGR2HSV)[:, :, 2]))
            except Exception as e:
                logger.warning(f"ROI处理异常: {str(e)}")
                continue

        if len(brightness) < 3:  # 至少需要3个有效亮度值
            return False

        try:
            fft = np.fft.fft(brightness)
            freqs = np.fft.fftfreq(len(brightness), d=1 / 30)
            mask = (np.abs(freqs) >= CONFIG['fire_validation']['flicker_min_hz']) & \
                   (np.abs(freqs) <= CONFIG['fire_validation']['flicker_max_hz'])
            energy_ratio = np.sum(np.abs(fft[mask])) / (np.sum(np.abs(fft)) + 1e-6)  # 避免除零
            return energy_ratio > CONFIG['fire_validation']['flicker_energy_ratio']
        except Exception as e:
            logger.error(f"FFT分析失败: {str(e)}")
            return False


if __name__ == '__main__':
    processors = []
    try:
        # 初始化视频处理器
        for area_id, video_path in VIDEO_PATHS.items():
            if not Path(video_path).exists():
                logger.error(f"视频文件不存在: {video_path}")
                continue

            logger.info(f"初始化视频处理器: {area_id}")
            processor = VideoProcessor(video_path, area_id)
            processors.append(processor)
            time.sleep(1)  # 避免同时启动多个模型导致显存溢出

        # 主循环
        while True:
            active_processors = [p for p in processors if p.active]
            if not active_processors:
                logger.warning("所有视频处理已完成")
                break

            time.sleep(5)
            logger.info(f"运行中... 活动处理器: {len(active_processors)}/{len(processors)}")

    except KeyboardInterrupt:
        logger.info("用户终止操作")
    except Exception as e:
        logger.error(f"主程序异常: {str(e)}", exc_info=True)
    finally:
        # 清理资源
        for p in processors:
            p.release()
        torch.cuda.empty_cache()
        logger.info("系统安全退出")